import logging

import numpy as np
import pandas as pd
import talib

logging.basicConfig(level=logging.INFO,
                    format='%(asctime)s | %(filename)s - {%(lineno)d} - [%(levelname)s] - %(message)s',
                    datefmt='%Y-%m-%d %H:%M:%S', encoding='utf-8')


class StockCalculate:

    def __init__(self, data_list):
        if data_list is not None and len(data_list) > 0:
            self.data_list = data_list
            self.df = pd.DataFrame(self.data_list,
                                   columns=["code", "date", "open", "close", "high", "low", "vol", "pchange"])
        else:
            logging.error("data_list is None or empty")
            raise ValueError("data_list is None or empty")

    def calculate(self):
        ## zl1,zl2 data list
        self.df["zz"] = (self.df["open"] + self.df["high"] + self.df["low"]) / 3

        self.df["zl1"] = talib.EMA(self.df["zz"].values, timeperiod=8)
        self.df["zl2"] = talib.EMA(self.df["zz"].values, timeperiod=55)
        ## ma5,ma10,ma20,ma30,ma60
        self.df["ma5"] = talib.SMA(self.df['close'].values.astype(float), timeperiod=5)
        self.df["ma10"] = talib.SMA(self.df['close'].values.astype(float), timeperiod=10)
        self.df["ma20"] = talib.SMA(self.df['close'].values.astype(float), timeperiod=20)
        self.df["ma30"] = talib.SMA(self.df['close'].values.astype(float), timeperiod=30)
        self.df["ma60"] = talib.SMA(self.df['close'].values.astype(float), timeperiod=60)
        ## vol5, vol10, vol20
        self.df["vol5"] = talib.SMA(self.df["vol"].values.astype(float), timeperiod=5)
        self.df["vol10"] = talib.SMA(self.df["vol"].values.astype(float), timeperiod=10)
        self.df["vol20"] = talib.SMA(self.df["vol"].values.astype(float), timeperiod=20)
        ## macd(dif,dea)
        self.df["dif"], self.df['dea'], self.df['bar'] = talib.MACD(self.df['close'].values, fastperiod=12,
                                                                    slowperiod=26, signalperiod=9)
        self.df["macd"] = (self.df["dif"] - self.df["dea"]) * 2
        ## close d1, d3, d5, d10
        self.df["c_d1"] = self.df["close"].shift(-1).rolling(1).max().fillna(0)
        self.df["c_d3"] = self.df["close"].shift(-3).rolling(3).max().fillna(0)
        self.df["c_d5"] = self.df["close"].shift(-5).rolling(5).max().fillna(0)
        self.df["c_d10"] = self.df["close"].shift(-10).rolling(10).max().fillna(0)
        ## high d1, d3, d5, d10
        self.df["h_d1"] = self.df["high"].shift(-1).rolling(1).max().fillna(0)
        self.df["h_d3"] = self.df["high"].shift(-3).rolling(3).max().fillna(0)
        self.df["h_d5"] = self.df["high"].shift(-5).rolling(5).max().fillna(0)
        self.df["h_d10"] = self.df["high"].shift(-10).rolling(10).max().fillna(0)
        ## low d1, d3, d5, d10
        self.df["l_d1"] = self.df["low"].shift(-1).rolling(1).max().fillna(0)
        self.df["l_d3"] = self.df["low"].shift(-3).rolling(3).max().fillna(0)
        self.df["l_d5"] = self.df["low"].shift(-5).rolling(5).max().fillna(0)
        self.df["l_d10"] = self.df["low"].shift(-10).rolling(10).max().fillna(0)
        # kdj
        self.df["kdj_k"], self.df["kdj_d"], self.df["kdj_j"] = self.calculate_kdj_tdx(self.df["high"], self.df["low"],
                                                                                      self.df["close"])
        ## RSI
        self.df["rsi_6"] = talib.RSI(self.df["close"].values, timeperiod=6)
        ## ATR
        self.df["atr_14"] = self.calculate_atr_tdx(self.df["high"].values, self.df["low"].values,
                                                   self.df["close"].values)

    def update(self):
        if self.df is not None and self.df.size > 0:
            return self.df
        else:
            return None

    def calculate_kdj_tdx(self, high, low, close, n=9):
        length = len(close)
        k = np.full(length, 50.0)
        d = np.full(length, 50.0)
        j = np.full(length, 0.0)
        rsv = np.full(length, 0.0)

        for i in range(n - 1, length):
            highest_high = np.max(high[i - n + 1:i + 1])
            lowest_low = np.min(low[i - n + 1:i + 1])

            if highest_high != lowest_low:
                rsv[i] = (close[i] - lowest_low) / (highest_high - lowest_low) * 100
            else:
                rsv[i] = 0
            if i == n - 1:
                k[i] = (2 / 3) * 50 + (1 / 3) * rsv[i]
                d[i] = (2 / 3) * 50 + (1 / 3) * k[i]
            else:
                k[i] = (2 / 3) * k[i - 1] + (1 / 3) * rsv[i]
                d[i] = (2 / 3) * d[i - 1] + (1 / 3) * k[i]
            j[i] = 3 * k[i] - 2 * d[i]

        return k, d, j

    def calculate_atr_tdx(self, high, low, close, period=14):
        tr1 = high - low
        tr2 = abs(high - np.roll(close, 1))
        mtr1 = np.maximum(tr1, tr2)
        tr3 = abs(low - np.roll(close, 1))
        tr = np.maximum(mtr1, tr3)
        return talib.SMA(tr, period)
